// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/compression/compression.h>
#include <pollux/dwio/common/output_stream.h>
#include <pollux/dwio/common/seekable_input_stream.h>
#include <pollux/dwio/common/compression/compression.h>
#include <pollux/dwio/common/compression/compression_buffer_pool.h>
#include <pollux/dwio/common/compression/paged_output_steam.h>
#include <pollux/dwio/dwrf/common/common.h>
#include <pollux/dwio/dwrf/common/config.h>
#include <pollux/dwio/dwrf/common/decryption.h>
#include <pollux/dwio/dwrf/common/encryption.h>

namespace kumo::pollux::dwrf {

using namespace dwio::common::compression;

constexpr uint8_t PAGE_HEADER_SIZE = 3;

inline CompressionOptions getDwrfOrcCompressionOptions(
    pollux::common::CompressionKind kind,
    uint32_t compressionThreshold,
    int32_t zlibCompressionLevel,
    int32_t zstdCompressionLevel) {
  CompressionOptions options{};
  options.compressionThreshold = compressionThreshold;

  if (kind == pollux::common::CompressionKind_ZLIB ||
      kind == pollux::common::CompressionKind_GZIP) {
    options.format.zlib.windowBits = Compressor::DWRF_ORC_ZLIB_WINDOW_BITS;
    options.format.zlib.compressionLevel = zlibCompressionLevel;
  } else if (kind == pollux::common::CompressionKind_ZSTD) {
    options.format.zstd.compressionLevel = zstdCompressionLevel;
  }
  return options;
}

/**
 * Create a compressor for the given compression kind.
 * @param kind The compression type to implement
 * @param bufferPool Pool for compression buffer
 * @param bufferHolder Buffer holder that handles buffer allocation and
 * collection
 * @param config The compression options to use
 */
inline std::unique_ptr<dwio::common::BufferedOutputStream> createCompressor(
    common::CompressionKind kind,
    CompressionBufferPool& bufferPool,
    dwio::common::DataBufferHolder& bufferHolder,
    const Config& config,
    const dwio::common::encryption::Encrypter* encrypter = nullptr) {
  CompressionOptions dwrfOrcCompressionOptions = getDwrfOrcCompressionOptions(
      kind,
      config.get(Config::COMPRESSION_THRESHOLD),
      config.get(Config::ZLIB_COMPRESSION_LEVEL),
      config.get(Config::ZSTD_COMPRESSION_LEVEL));
  auto compressor = createCompressor(kind, dwrfOrcCompressionOptions);
  if (!compressor) {
    if (!encrypter && kind == common::CompressionKind::CompressionKind_NONE) {
      return std::make_unique<dwio::common::BufferedOutputStream>(bufferHolder);
    }
  }
  return std::make_unique<PagedOutputStream>(
      bufferPool,
      bufferHolder,
      dwrfOrcCompressionOptions.compressionThreshold,
      PAGE_HEADER_SIZE,
      std::move(compressor),
      encrypter);
}

inline CompressionOptions getDwrfOrcDecompressionOptions(
    common::CompressionKind kind) {
  CompressionOptions options{};
  if (kind == common::CompressionKind_ZLIB ||
      kind == common::CompressionKind_GZIP) {
    options.format.zlib.windowBits = Compressor::DWRF_ORC_ZLIB_WINDOW_BITS;
  } else if (
      kind == common::CompressionKind_LZ4 ||
      kind == common::CompressionKind_LZO) {
    options.format.lz4_lzo.isHadoopFrameFormat = false;
  }
  return options;
}

/**
 * Create a decompressor for the given compression kind.
 * @param kind The compression type to implement
 * @param input The input stream that is the underlying source
 * @param bufferSize The maximum size of the buffer
 * @param pool The memory pool
 */
inline std::unique_ptr<dwio::common::SeekableInputStream> createDecompressor(
    kumo::pollux::common::CompressionKind kind,
    std::unique_ptr<dwio::common::SeekableInputStream> input,
    uint64_t bufferSize,
    memory::MemoryPool& pool,
    const std::string& streamDebugInfo,
    const dwio::common::encryption::Decrypter* decryptr = nullptr) {
  const CompressionOptions& options = getDwrfOrcDecompressionOptions(kind);
  return createDecompressor(
      kind,
      std::move(input),
      bufferSize,
      pool,
      options,
      streamDebugInfo,
      decryptr);
}

} // namespace kumo::pollux::dwrf
